subprocess_venv: two-stage install to fix torch/torchaudio CUDA-tag s…#518
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements a two-stage installation process for virtual environments using uv to resolve an ABI mismatch issue caused by index priority. Stage 1 installs torch and torchaudio from a specific CUDA index, while Stage 2 installs the remaining requirements from PyPI. Feedback suggests filtering torch-related packages from the Stage 2 requirements to prevent re-resolution and improving the _torch_install_specs function to handle multiple version constraints for the same package.
| # via default PyPI. ``torchaudio`` is already installed in Stage 1 | ||
| # so it isn't re-added here; safetensors + numpy are pure-Python | ||
| # and pull cleanly from PyPI on every supported platform. | ||
| stage_two_reqs = [*requirements, "safetensors", "numpy"] |
There was a problem hiding this comment.
To ensure that torch and torchaudio are never re-resolved against the default PyPI index in Stage 2, it is safer to filter them out of the requirements list passed to the second subprocess.run. The PR description notes that re-listing them without an index flag could lead uv to replace the matched wheels with PyPI's tagless ones. While this was addressed for the explicit torchaudio append, backends that include these packages in their own requirements (like NeMo or Canary) still re-list them in Stage 2, which might trigger the behavior you are trying to avoid.
| stage_two_reqs = [*requirements, "safetensors", "numpy"] | |
| stage_two_reqs = [r for r in requirements if _spec_pkg_name(r) not in _TORCH_PKG_NAMES] + ["safetensors", "numpy"] |
| pinned: dict[str, str] = {} | ||
| for spec in requirements: | ||
| name = _spec_pkg_name(spec) | ||
| if name in _TORCH_PKG_NAMES: | ||
| pinned[name] = spec | ||
| return [pinned.get(name, name) for name in ("torch", "torchaudio")] |
There was a problem hiding this comment.
The current implementation of _torch_install_specs only keeps the last constraint found in requirements for each package because it uses a dictionary. If a backend specifies multiple constraints for torch (e.g., ["torch>=2.8", "torch<2.9"]), only one will be passed to the Stage 1 install. This could lead to Stage 1 installing a version that doesn't satisfy the other constraint, which Stage 2 might then try to "fix" by pulling from PyPI, defeating the routing logic. It is safer to collect and return all matching specs.
out = []
found = set()
for spec in requirements:
name = _spec_pkg_name(spec)
if name in _TORCH_PKG_NAMES:
out.append(spec)
found.add(name)
for name in ("torch", "torchaudio"):
if name not in found:
out.append(name)
return out
|
the only consideration is that if a subprocess venv does not require torchaudio (or relies only on torchcodec), this forces it. especially as this approach is meant not just for legacy environments but also when the environment needs something from simply a github repo. |
3f0cb35 to
bc98113
Compare
|
Added a
So a future backend that consumes a pure-Python GitHub repo can declare: ensure_venv("my-backend", ["some-repo @ git+https://..."], python_version="3.12", requires_torch=False)
and pay none of the multi-GB CUDA wheel cost. Test coverage in test_requires_torch_false_skips_probe_and_stage_one asserts the probe is never invoked and the single install
argv has neither --index-url nor torchaudio.
Worth flagging that the current all_reqs = [..., "safetensors", "numpy", "torchaudio"] line predates this PR — the unconditional torchaudio append goes back to the original
ensure_venv implementation. This patch is the first place that lets a caller opt out of it. If we want to be more aggressive (e.g. drop torchaudio from the default IPC list
and only add it when a backend explicitly needs it for FLAC encoding), that's a separate refactor — happy to do it as a follow-up if you'd rather change the default than gate
it behind a flag. |
|
but torch doesn't always require torchaudio, does it? i thought it was removed as a requirement. so the only time this comes in is when some other thing requires it. |
|
Right — torch doesn't pull torchaudio (the dep direction is the other way). The unconditional So your point is sharper than what Honest fix is to drop the unconditional append from
Two ways to land that:
Mild preference for (2) since each PR's effect is then isolated and bisectable, but (1) is straightforward if you'd rather not see two PRs back-to-back. Which way? |
…equirements PR #516 routes the install through ``--index-url <cuda> --extra-index-url https://pypi.org/simple`` so non-torch packages can still resolve from PyPI. But uv treats ``--extra-index-url`` as having higher priority than ``--index-url`` (opposite of pip), so PyPI wins for every package — torch and torchaudio included. On hosts where PyPI ships those with mismatched ``+cu`` local-version tags (currently ``torch==X+cu129`` vs ``torchaudio==X`` with no tag, internally cu128), the resulting venv reproduces the same ABI mismatch this PR is meant to fix: RuntimeError: Detected that PyTorch and TorchAudio were compiled with different CUDA versions. PyTorch has CUDA version 12.9 whereas TorchAudio has CUDA version 12.8. Reproduced on MIT ORCD with the env override set (``SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128``) — override URL was honored in the marker but uv still resolved both wheels from PyPI. So the override path was non-functional too. Fix: split the install into two stages and decide whether Stage 1 fires from the caller's ``requirements`` itself. Stage 1 — ``uv pip install --index-url <chosen>`` listing every ``torch`` / ``torchaudio`` spec found in ``requirements``, with NO ``--extra-index-url``. The chosen CUDA index is unambiguously primary; both wheels and their ``nvidia-cuda-runtime-cu12`` transitives come from the same toolchain. Multiple constraints for the same package (``["torch>=2.8", "torch<2.9"]``) are all forwarded so uv combines them at resolve time. Stage 2 — ``uv pip install <rest of requirements> safetensors numpy`` with no index flags. Backend-pinned torch / torchaudio specs are filtered out so uv can't be tempted to re-resolve them from PyPI against the matched ``+cu128`` wheels installed in Stage 1. Drop the unconditional ``torchaudio`` IPC append, in two ways: 1. ``_torch_install_specs`` no longer pads its return with bare ``torch`` / ``torchaudio`` names. If neither is in ``requirements``, it returns ``[]``. 2. ``ensure_venv`` treats ``_torch_install_specs(requirements) == []`` as the trigger to skip the probe entirely and run a single torch-free install pass against default PyPI. No ``nvidia-smi`` shellout, no ``torchaudio`` force-appended. This addresses @satra's review concern that the previous default forced ``torchaudio`` into every subprocess venv, even backends that don't use it. With the change, ``yamnet`` and ``continuous-ser`` — both of which read audio via ``soundfile`` in their worker scripts and never import torchaudio — get ~200 MB leaner venvs. Backends that genuinely need ``torch`` / ``torchaudio`` (including via a transitive dep) MUST pin them explicitly so Stage 1 routes them through the matched CUDA index; this is also why ``qwen.py``'s ``_REQUIREMENTS`` now names them directly even though ``qwen-asr==0.0.6`` would otherwise pull them transitively — relying on the transitive would let Stage 2's PyPI resolution split the local-version tags. The 7 backends that already pin ``torchaudio`` explicitly (canary, nemo, ppgs, sparc, coqui, s3prl) are unaffected. The earlier ``requires_torch=False`` parameter from a previous round of this review (an explicit opt-out flag) is removed in favor of the auto-detection — it's structurally equivalent and the explicit flag becomes redundant noise. Behavioral guarantees preserved from PR #516: marker schema (extended, not changed), cache-hit fast path, env override (``SENSELAB_TORCH_INDEX_URL``), ``SenselabCudaCompatibilityError`` wrapping on wheel-not-found errors (now scoped to Stage 1 where it semantically belongs), pass-through of unrelated ``CalledProcessError``, half-built-venv cleanup on failure, host-CUDA probe diagnostic surface, cross-backend routing. Tests: 34 in total covering (1) Stage 1 names only ``--index-url`` with torch + torchaudio pinned per requirements; (2) Stage 2 carries no index flags + IPC deps; (3) Stage 2 filters torch / torchaudio specs from caller's requirements; (4) multiple constraints for the same package all forward through Stage 1 verbatim; (5) requirements without torch / torchaudio skip the probe and Stage 1 entirely; (6) yamnet- style requirements get NO torchaudio in their install argv; (7) switching a venv name from torch-free to torch-using correctly invalidates the cache and rebuilds. The cross-backend regression test asserts both that Stage 1 routes through the chosen index AND that Stage 2 contains no torch / torchaudio specs at all. Review feedback addressed: - @satra (PR #518 review): drop the unconditional ``torchaudio`` IPC append; auto-detect torch routing from the caller's ``requirements`` instead of forcing every venv to pay for it. Adds explicit ``torch`` + ``torchaudio`` to qwen's ``_REQUIREMENTS`` so its transitive-dep resolution still flows through the CUDA index. - gemini-code-assist #1: filter torch / torchaudio specs from Stage 2. - gemini-code-assist #2: list-based ``_torch_install_specs`` keeps every matching constraint instead of clobbering through a dict. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
bc98113 to
80db255
Compare
2c0d351
into
20260512-204619-fix-canary-cuda-conflict
2 participants
Summary
Follow-up to #516. The
--index-url <cuda> --extra-index-url pypipattern doesn't actually routetorch/torchaudiothrough the chosen PyTorch wheel index in uv.Splitting the install into two stages (CUDA-index for torch+torchaudio, then default PyPI for everything else) makes the routing real, eliminates the cu128/cu129 ABI mismatch
this PR was meant to fix, and keeps every behavioral guarantee #516 introduced.
Why
While testing #516 on ORCD, the canary venv consistently rebuilt into the same broken state #516 is trying to prevent:
Reproduced even with
SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128explicitly set:The marker reported
torch_index.url == https://download.pytorch.org/whl/cu128(override was honored at the routing-decision layer), but the install argv still landed bothwheels from PyPI.
Root cause
uv treats
--extra-index-urlas taking precedence over--index-url(opposite of pip). Quoting uv's own docs: "uv treats--extra-index-urlas a list of extra indexesthat take precedence over
--index-url." So in #516's install command:PyPI is consulted first for every package. PyPI has both
torchandtorchaudio, so the CUDA index is effectively never used for them. The two wheels come from PyPI withmismatched local-version tags, and the env override path is non-functional too.
Verified directly: running
uv pip install --index-url https://download.pytorch.org/whl/cu128 --force-reinstall --no-deps torch torchaudio(no--extra-index-url) correctlyinstalled matched
+cu128wheels into the same venv.The fix
Split the install into two stages.
Stage 1 —
uv pip install --index-url <chosen> torch torchaudiowith no--extra-index-url. The chosen CUDA index is unambiguously primary; both wheels and theirnvidia-cuda-runtime-cu12transitives come from the same toolchain. Any version pin in the caller'srequirements(e.g.torch>=2.8,<2.9) flows through verbatim; backendsthat don't pin torch (e.g. qwen) fall back to bare names.
Stage 2 —
uv pip install <rest of requirements> safetensors numpywith no index flags. uv resolves against default PyPI, seestorch/torchaudioalready installedand satisfying any pin, and leaves them alone. The PyTorch index now governs only the two packages it's designed for, and stale wheels on the CUDA index for utilities like
setuptools / pyarrow stay out of the picture.
torchaudiois no longer re-appended in Stage 2 — it's installed in Stage 1, and re-listing it without an index flag would let uv consider replacing the matched wheel withPyPI's tagless one.
Behavioral guarantees preserved from #516
torch_indexfield withtag/url/source).requirements+torch_index.url).SENSELAB_TORCH_INDEX_URL) — flows into Stage 1 unchanged.SenselabCudaCompatibilityErrorwrapping on wheel-not-found errors, now scoped to Stage 1 where it semantically belongs (Stage 2 hits default PyPI, so a failure thereisn't a CUDA-compat surface).
CalledProcessErrors.ensure_venv).Tests
The 14 existing
subprocess_venv_test.pytests are kept; 4 needed call-count or argv-shape updates for the new two-stage form, and the single install-argv test was splitinto three more specific ones:
test_stage_one_pins_torch_and_torchaudio_to_chosen_index— Stage 1 names only--index-url, no--extra-index-url, and includes the torch + torchaudio specs fromrequirements.test_stage_one_uses_bare_names_when_requirements_omit_torch— qwen-style backend (no explicit torch pin) still gets torch + torchaudio installed in Stage 1 so transitivesdon't pull them in later from PyPI.
test_stage_two_uses_default_pypi_and_includes_ipc_deps— Stage 2 carries neither--index-urlnor--extra-index-url, includessafetensors+numpy, and does NOTre-list
torchaudio.The cross-backend regression test (
test_all_three_subprocess_backends_route_through_same_torch_index) now asserts both that Stage 1 routes through the chosen index AND thatStage 2 does not — guarding against any future re-introduction of
--extra-index-urlon the torch install. TheSenselabCudaCompatibilityErrorwrapping test and theunrelated-error pass-through test were unchanged (they fire on the first
uv pip installcall, which is now Stage 1 — semantically correct).All 30 unit/integration tests pass:
Out-of-band verification
Reproduced the bug, applied this fix, and verified end-to-end on an ORCD compute node:
~/.cache/senselab/venvs/nemo-canary-qwento force a fresh build.cpu.torch-2.8.0+cpuandtorchaudio-2.8.0+cpu(clean CPU path — no PyPI fallback leaking cu128/cu129).RuntimeError. CPU path proven functional.Inside a
srun --gres=gpu:1allocation the probe pickscu128from the static map and Stage 1 installs matched+cu128wheels — same outcome, GPU-capable.Test plan
uv run pytest src/tests/utils/subprocess_venv_test.py src/tests/utils/cuda_probe_test.py --noconftest— all 30 tests pass.SENSELAB_VENV_CACHEredirected to a throwaway dir; verified torch + torchaudio wheels carry matching local-version tags after theinstall.
temp/asr.pyruns canary end-to-end against a real audio clip; noRuntimeError, transcript returned.